from asgiref.sync import async_to_sync
from channels.generic.websocket import WebsocketConsumer
import json
import base64
from django.http.request import QueryDict
import paramiko
from threading import Thread
import os
from io import StringIO
from django.conf import settings


# 解析密钥文件
def get_key_obj(pkeyobj, pkey_file=None, pkey_obj=None, password=None):
    if pkey_file:
        with open(pkey_file) as fo:
            try:
                pkey = pkeyobj.from_private_key(fo, password=password)
                return pkey
            except:
                pass
    else:
        try:
            pkey = pkeyobj.from_private_key(pkey_obj, password=password)
            return pkey
        except:
            pkey_obj.seek(0)


"""
   桥接WebSocket和ssh
"""


class SSHBridge(object):

    def __init__(self, websocket):
        self.websocket = websocket
        self.ssh_channel = None

    def connect(self, host, user, pwd=None, key=None, port=22, timeout=6, term='xterm', pty_width=80, pty_height=24):
        # 建立ssh连接
        ssh_client = paramiko.SSHClient()
        ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            if key:
                # 密钥方式认证
                pkey = get_key_obj(paramiko.RSAKey, pkey_obj=key, password=pwd) or \
                       get_key_obj(paramiko.DSSKey, pkey_obj=key, password=pwd) or \
                       get_key_obj(paramiko.ECDSAKey, pkey_obj=key, password=pwd) or \
                       get_key_obj(paramiko.Ed25519Key, pkey_obj=key, password=pwd)
                ssh_client.connect(username=user, hostname=host, port=port, pkey=pkey, timeout=timeout)
            else:
                ssh_client.connect(hostname=host, port=port, username=user, password=pwd, timeout=timeout)
        except Exception as e:
            message = json.dumps({'flag': 'fail', 'message': str(e)})
            self.websocket.send_message_or_team(message)
            return

        transport = ssh_client.get_transport()
        # 打开一个通道
        self.ssh_channel = transport.open_session()
        # 获取一个终端
        self.ssh_channel.get_pty(term=term, width=pty_width, height=pty_height)
        # 激活终端，这样就可以登录到终端了，就和我们用类似于xshell登录系统一样
        self.ssh_channel.invoke_shell()

        # 获取ssh连接主机后的返回内容，例如Linux，会显示上次登录等信息，把这些信息通过WebSocket显示到Web终端。
        # 连接建立一次，之后交互数据不会再进入该方法
        for i in range(2):
            # 这里因为可能出现编码问题
            try:
                recv = self.ssh_channel.recv(1024).decode('utf-8')
            except Exception as e:
                recv = self.ssh_channel.recv(1024).decode('gbk')
            message = json.dumps({'flag': 'msg', 'message': recv})
            self.websocket.send_message_or_team(message)

    def close(self):
        message = {'flag': 0, 'message': '关闭WebSocket和SSH连接'}
        # 向WebSocket发送一个关闭消息
        self.websocket.send_message_or_team(json.dumps(message))

        try:
            # 关闭ssh通道
            self.ssh_channel.close()
            # 关闭WebSocket连接
            self.websocket.close()
        except BaseException as e:
            print('关闭WebSocket和SSH连接产生异常：', e)
            pass

    def _ws_to_ssh(self, data):
        # 尝试发送数据到ssh通道，产生异常则关闭所有连接
        try:
            self.ssh_channel.send(data)
        except OSError:
            self.close()

    def _ssh_to_ws(self):
        try:
            while not self.ssh_channel.exit_status_ready():
                try:
                    data = self.ssh_channel.recv(1024).decode('utf-8')
                except Exception as e:
                    data = self.ssh_channel.recv(1024).decode('gbk')
                if len(data) != 0:
                    message = {'flag': 'msg', 'message': data}
                    self.websocket.send_message_or_team(json.dumps(message))
                else:
                    break
        except Exception as e:
            self.close()

    def shell(self, data):
        # 开启两个线程，一个接收消息一个发送消息
        Thread(target=self._ws_to_ssh, args=(data,)).start()
        Thread(target=self._ssh_to_ws).start()

    def resize_pty(self, cols, rows):
        # 重新改变虚拟终端长和宽
        self.ssh_channel.resize_pty(width=cols, height=rows)


class WebsshConsumer(WebsocketConsumer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.host_id = ''
        self.simple_user = ''
        self.is_team = False
        self.team_name = ''
        self.ssh = None

    def connect(self):
        # 建立WebSocket连接，并实例化SSHBridge类，在这个对象中建立SSH连接
        self.host_id = self.scope['url_route']['kwargs'].get('host_id')
        self.accept()

        # WebSocket连接成功后，连接ssh
        query_string = self.scope.get('query_string')
        ws_args = QueryDict(query_string=query_string, encoding='utf-8')
        width = ws_args.get('width')
        height = ws_args.get('height')
        width = int(width)
        height = int(height)  # ssh连接要求int类型：required argument is an integer
        user = ws_args.get('user')
        host = ws_args.get('host')
        port = ws_args.get('port')
        port = int(port)
        auth = ws_args.get('auth')
        pwd = ws_args.get('pwd')
        if pwd:
            pwd = base64.b64decode(pwd).decode('utf-8')
        sshkey_filename = ws_args.get('sshkey_filename')

        ssh_connect_dict = {
            'host': host,
            'user': user,
            'port': port,
            'timeout': 30,
            'pty_width': width,
            'pty_height': height,
            'pwd': pwd
        }

        if auth == 'key':
            sshkey_file = os.path.join(settings.MEDIA_ROOT, 'sshkey', sshkey_filename)
            if not os.path.exists(sshkey_file):
                self.send(json.dumps({'flag': 'error', 'message': '密钥文件不存在'}))

            else:
                try:
                    f = open(sshkey_file, 'r', encoding='utf-8')
                    key = f.read()
                    string_io = StringIO()
                    string_io.write(key)
                    string_io.flush()
                    string_io.seek(0)
                    ssh_connect_dict['key'] = string_io

                    os.remove(sshkey_file)  # 用完之后删除key文件
                except BaseException as e:
                    pass

        # 建立SSH连接
        self.ssh = SSHBridge(websocket=self)
        self.ssh.connect(**ssh_connect_dict)

    def disconnect(self, close_code):
        # 断开连接
        try:
            self.ssh.close()
        except BaseException as e:
            pass

    def receive(self, text_data=None, bytes_data=None):
        # 从WebSocket中接收消息
        text_data = json.loads(text_data)  # json字符串转字典
        if type(text_data) == dict:
            if text_data.get('flag') == 'entered_key':
                data = text_data.get('entered_key', '')  # 获取前端传过来输入的按键值，并传递给shell
                self.ssh.shell(data=data)
            else:
                cols = text_data['cols']
                rows = text_data['rows']
                # 改变通道中终端大小
                self.ssh.resize_pty(cols=cols, rows=rows)
        else:
            pass

    def send_message_or_team(self, message):
        self.send(message)

    def team_message(self, event):
        message = event['message']
        self.send(message)
